{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bb5e5a3a",
   "metadata": {},
   "source": [
    "## 自注意力机制\n",
    "$$\n",
    "Attention(Q, K, V) = softmax(\\frac{QK^T}{\\sqrt{d_k}})V\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7337802a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 10, 64])\n",
      "torch.Size([32, 10, 64])\n"
     ]
    }
   ],
   "source": [
    "# 自注意力机制实现\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import math\n",
    "\n",
    "\n",
    "class SelfAttention(nn.Module):\n",
    "    def __init__(self, dim_q, dim_k, dim_v):\n",
    "        super(SelfAttention, self).__init__()\n",
    "        self.dim_q = dim_q\n",
    "        self.dim_k = dim_k\n",
    "        self.dim_v = dim_v\n",
    "        # 定义线性层变化\n",
    "        self.linear_q = nn.Linear(dim_q, dim_k, bias=False)\n",
    "        self.linear_k = nn.Linear(dim_q, dim_k, bias=False)\n",
    "        self.linear_v = nn.Linear(dim_q, dim_v, bias=False)\n",
    "\n",
    "        # 缩放点积注意力的关键部分，用于防止点积过大导致softmax梯度消失\n",
    "        self._norm_fact = 1/ math.sqrt(dim_k)\n",
    "\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # 输入形状 [batch_size, sequence_length, input_dim]\n",
    "        batch, n, dim_q = x.shape\n",
    "\n",
    "        # 生成查询Q , 键K , 值V\n",
    "        q = self.linear_q(x)\n",
    "        k = self.linear_k(x)\n",
    "        v = self.linear_v(x)\n",
    "\n",
    "        # 计算注意力分数\n",
    "        dist = torch.bmm(q, k.transpose(1, 2)) * self._norm_fact\n",
    "        dist = F.softmax(dist, dim=-1)\n",
    "\n",
    "        # 加权求和\n",
    "        att = torch.bmm(dist, v)\n",
    "\n",
    "        return att\n",
    "\n",
    "batch_size = 32\n",
    "seq_len = 10\n",
    "input_dim = 64\n",
    "\n",
    "self_attn = SelfAttention(dim_q=64, dim_k=32, dim_v=64)\n",
    "x = torch.randn(batch_size, seq_len, input_dim)\n",
    "out = self_attn(x)\n",
    "print(x.shape)\n",
    "print(out.shape)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5107559b",
   "metadata": {},
   "source": [
    "## 通道注意力机制"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1e7a48ef",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 64, 32, 32])\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "class ChannelAttention(nn.Module):\n",
    "    def __init__(self, in_channels, reduction_ratio=16):\n",
    "        super(ChannelAttention, self).__init__()\n",
    "        \"\"\"\n",
    "        reduction_ratio: 缩减比率, 用于降维\n",
    "        \"\"\"\n",
    "        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n",
    "        self.max_pool = nn.AdaptiveMaxPool2d(1)\n",
    "\n",
    "        self.fc = nn.Sequential(\n",
    "            nn.Linear(in_channels, in_channels // reduction_ratio, bias=False),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(in_channels // reduction_ratio, in_channels, bias=False)\n",
    "        )\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, x: torch.Tensor):\n",
    "        b, c, h, w = x.shape\n",
    "        # view(b, c)重塑为 [32, 64, 1, 1] -> [32, 64]\n",
    "        avg_out = self.fc(self.avg_pool(x).view(b, c))\n",
    "        max_out = self.fc(self.max_pool(x).view(b, c))\n",
    "\n",
    "        out = avg_out + max_out\n",
    "        # 压缩到[0, 1]的范围内\n",
    "        channel_weights = self.sigmoid(out).view(b, c, 1, 1)\n",
    "        # expand_as 将权重[32, 64, 1, 1] 扩充到 [32, 64, 32, 32]\n",
    "        # 逐个元素相乘\n",
    "        return x * channel_weights.expand_as(x)\n",
    "    \n",
    "bt_size, channels, h, w = 32, 64, 32, 32\n",
    "x = torch.randn(bt_size, channels, h, w)\n",
    "# print(type(x))\n",
    "channel_att = ChannelAttention(channels)\n",
    "out = channel_att(x)\n",
    "print(out.shape)\n",
    "\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
